#include "simscrub.hpp"

#ifdef LDPC
const static char* mat_dir_H1 = (char*) "matrices/H4376";  // the directory of the parity check matrix H1
//const static char* mat_dir_H1 = (char*) "matrices/H8000";  // the directory of the parity check matrix H1
#endif
const static char* mat_dir_H2 = (char*) "matrices/H3584.443";  // the directory of the parity check matrix H2

// global counters and switches
double average_iter_num;
double average_error_num;
double average_correct_blk_num;
double average_error_blk_num;
double average_undet_blk_num;
double average_det_blk_num;
double average_final_ber;
double average_mem_bit_num;       // number of bits memoized before each call to ECC 
double average_cleanup_bit_num;   // number of bits get cleaned up 
double num_seg_scrubbed;          // number of segments which are scrubbed.
double scrub_gap;
double scrub_gap_naive;
double total_time;
int trial_num;
vector<double> bit_error_rates;
int scrubbing_selection_switch;

double trial_num_meta;
double average_ber_before_decode_meta;
double average_iter_num_meta;
double average_error_num_meta;
double average_correct_blk_num_meta;
double average_error_blk_num_meta;
double average_compressed_metadata_size;

long BCHk;                                      
GF2X BCHgenerator;                              

// print the performance result.
void print_statistics(){
  cout << endl << "----------------------------------------------" << endl;
  cout << "Decoding Performance (Ave.)" << endl;
  cout << "==============================================" << endl;
  cout << fixed << setprecision(6);
  cout << "1.  Final error#: " << average_error_num << endl;
  cout << "2.  Final BER: " << average_final_ber << endl;
  cout << "3.  Correct blk#: " << average_correct_blk_num << endl;
  cout << "4.  Error blk#: " << average_error_blk_num << endl;
#ifdef LDPC
  cout << "5.  Iteration: " << average_iter_num << endl;
  cout << "6.  Undetected blk#: " << average_undet_blk_num << endl;
  cout << "7.  Detected blk#: " << average_det_blk_num << endl;
  cout << "8.  Success rate: " << average_correct_blk_num/(average_correct_blk_num +average_undet_blk_num + average_det_blk_num) << endl; 
#endif

  if (scrubbing_selection_switch == 1){
    cout << "9.  Total seg scrubbed: " << num_seg_scrubbed << endl;
    cout << "10. Memoized bit#: " << average_mem_bit_num << endl;
    cout << "11. Refreshed bit#: " << average_cleanup_bit_num << endl;
    cout << "12. BER before metadata decoding: " << average_ber_before_decode_meta << endl;
    cout << "13. BER of metadata decoding: " << average_error_num_meta / ((double) N2) << endl;
    cout << "14. Iteration of metadata decoding: " << average_iter_num_meta << endl;
    cout << "15. Correct blk# of metadata decoding: " << average_correct_blk_num_meta << endl;
    cout << "16. Error blk# of metadata decoding: " << average_error_blk_num_meta << endl;
    cout << "17. Size of compressed metadata: " << average_compressed_metadata_size << endl;
  }
  cout << "----------------------------------------------" << endl;
}

// print out simulation information

void print_info(){
  cout << endl;
  cout << "----------------------------------------------" << endl;
  cout << "Experimental setup" << endl;
  cout << "==============================================" << endl;
#ifdef LDPC
  cout << "C1 Block length: " << N << endl;
  cout << "C1 Parity check node#: " << M << endl;
  cout << "C1 Code rate: " << (double) K/N << endl;
#endif
#ifdef BCH
  cout << "C1 Block length:" << BCHn << endl;
  cout << "Message bits# K:" << BCHk << endl;
  cout << "C1 Code rate:" << static_cast<double>(BCHk)/static_cast<double>(BCHn) << endl; 
  cout << "C1 Parity check bits# M:" << BCHn - BCHk << endl;
  cout << "p:" << 2 << endl;
  cout << "m:" << BCHm << endl; // don't confuse this m with M which is the number of parity check bits
  cout << "t:" << BCHt << endl;
  cout << "Generator: " << BCHgenerator << endl;
  cout << "Generator degree: " << deg(BCHgenerator) << endl;
#endif
  cout << "C2 Block length: " << N2 << endl;
  cout << "C2 Parity check node#: " << M2 << endl;
  cout << "C2 Code rate: " << (double) K2/N2 << endl;
  cout << "LDPC decoding iteration limit: " << MAX_ITER << endl;
  cout << "Total time (s): " << total_time << endl;
  cout << "Scrub gap (s): ";
  if (scrubbing_selection_switch == 1){
    cout << scrub_gap << endl;
    cout << "Scheme: multi-phase" << endl;
  }
  else{
    cout << scrub_gap_naive << endl;
    cout << "Scheme: naive" << endl;
  }
  cout << "----------------------------------------------" << endl << endl;
}

// read in the parity-check matrix and store the matrix in a 1-d
// array.
void read_matrix_H(mat_GF2& H, const char* dir, const int m, const int n){
  int b;
  H.SetDims(m, n);  // M * N matrix
  ifstream infile(dir);
  if (!infile) {
    cout << "The following file can not be opened:" << dir << endl;
    exit(EXIT_FAILURE);
  }
  for(long i = 1; i <= m; i++){
    for(long j = 1; j <= n; j++){
      infile >> b;
      H(i, j) = to_GF2(b);
    }
  }
}


// switch = 1: staged scrubbing scheme.
// switch = 2: naive scrubbing scheme.
void simulate_scrub_ldpc(const mat_GF2& H, const mat_GF2& G){
  vec_GF2 codeword;
  mat_GF2 H2 = mat_GF2();
  mat_GF2 G2 = mat_GF2();
  // for multi-phase scheme, we need to read in H2 for protecting the
  // metadata.
  if (scrubbing_selection_switch == 1){
    read_matrix_H(H2, mat_dir_H2, M2, N2);
    LDPCEncoder::compute_G(G2, H2);
  }

  print_info();

#ifdef VERBOSE
  if (scrubbing_selection_switch == 1)
    cout << "Using multi-phase scrubbing" << endl;
  else 
    cout << "Using naive scrubbing" << endl;
#endif

  for(long i = 1; i <= trial_num; i++){
    LDPCEncoder::generate_random_codeword(codeword, G);

#ifdef VERBOSE
      cout << "Start the " << i << "-th round of experiment" << endl;
#endif
      if (scrubbing_selection_switch == 1)
#ifndef PR1
	Algorithms::multiphase_scrubbing_ldpc(codeword, H, H2, G2);
#else
        Algorithms::multiphase_scrubbing_ldpc(codeword, H);
#endif
      else
	Algorithms::naive_scrubbing_ldpc(codeword, H);
  }
}

// FIXME: the performance measurement part is not implemented 
// the function for simulating memory scrubbing.
// switch = 1: staged scrubbing scheme.
// switch = 2: naive scrubbing scheme.
void simulate_scrub_bch(){
  vec_GF2 codeword;
  mat_GF2 H2 = mat_GF2();
  mat_GF2 G2 = mat_GF2();

  // for multi-phase scheme, we need to read in H2 for protecting the
  // metadata.
  if (scrubbing_selection_switch == 1){
    read_matrix_H(H2, mat_dir_H2, M2, N2);
    LDPCEncoder::compute_G(G2, H2);
  }

  GF2X genpol;                                  
  BCHEncoder::compute_generator_polynomial(genpol);
  BCHgenerator = genpol;
  BCHk = BCHn - deg(BCHgenerator);
  print_info();

#ifdef VERBOSE
  if (scrubbing_selection_switch == 1)
    cout << "Using multi-phase scrubbing" << endl;
  else 
    cout << "Using naive scrubbing" << endl;
#endif

  for(long i = 1; i <= trial_num; i++){
    BCHEncoder::generate_random_codeword(codeword, genpol);
#ifdef VERBOSE
      cout << "Start the " << i << "-th round of experiment" << endl;
#endif
      if (scrubbing_selection_switch == 1){
	Algorithms::multiphase_scrubbing_bch(codeword, H2, G2);
      }
      else
	Algorithms::naive_scrubbing_bch(codeword);
  }
}



// initialize the global performance counters
void init_performance_counters(){
  //fill(bit_error_rates.begin(), bit_error_rates.end(), 0.0);
  average_iter_num = 0;
  average_error_num = 0;
  average_correct_blk_num = 0;
  average_error_blk_num = 0;
  average_undet_blk_num = 0;
  average_det_blk_num = 0;
  average_final_ber = 0;
  average_mem_bit_num = 0;      
  average_cleanup_bit_num = 0;  
  num_seg_scrubbed = 0;

  trial_num_meta = 0;
  average_ber_before_decode_meta = 0;
  average_iter_num_meta = 0;
  average_error_num_meta = 0;
  average_correct_blk_num_meta = 0;
  average_error_blk_num_meta = 0;
  average_compressed_metadata_size = 0;
}

// average the performance counters
void compute_average_performance(){
  for(long i = 0; i < (long) bit_error_rates.size(); i++){
    bit_error_rates[i] /= trial_num;
  }
  if (scrubbing_selection_switch == 1){
    average_iter_num /= (trial_num * (total_time/scrub_gap)) ;
    num_seg_scrubbed /= trial_num;
    average_mem_bit_num /= (num_seg_scrubbed * trial_num); 
    average_cleanup_bit_num /= (num_seg_scrubbed * trial_num);

    average_ber_before_decode_meta /= trial_num_meta;
    average_iter_num_meta /= trial_num_meta;
    average_error_num_meta /= trial_num_meta;
    average_correct_blk_num_meta /= trial_num_meta;
    average_error_blk_num_meta /= trial_num_meta;
    average_compressed_metadata_size /= trial_num_meta;
  }
  else
    average_iter_num /= (trial_num * (total_time/scrub_gap_naive)) ;
  average_error_num /= trial_num;
  average_correct_blk_num /= trial_num;
  average_error_blk_num /= trial_num;
  average_undet_blk_num /= trial_num;
  average_det_blk_num /= trial_num;
#ifdef LDPC
  average_final_ber = average_error_num / N;
#endif

#ifdef BCH
  average_final_ber = average_error_num / BCHn;
#endif
}

int get_directory_name(){
  int h;
  if (scrubbing_selection_switch == 1){
#ifdef LDPC
    h = (scrub_gap * (N / (BITS_PER_CELL * CELL_PER_SEG)))/3600;
#endif
#ifdef BCH
    h = (scrub_gap * ((BCHn + 1) / (BITS_PER_CELL * CELL_PER_SEG)))/3600;
    std::cout << "scrub gap:" << scrub_gap << std::endl;
    if (scrub_gap == 450.0)
      h = 5;
#endif
  }
  else{
    h = scrub_gap_naive / 3600;
  }
  return h;
}


// generate a file name to store the performance result.
const string get_file_name(){
  stringstream ss;
  ss << "output/" << get_directory_name() << "/";
  if (scrubbing_selection_switch == 1){
#ifdef LDPC
    ss << "LS";
#endif 
#ifdef BCH
    ss << "BS";
#endif
    }
  else{
#ifdef LDPC
    ss << "LN";
#endif 
#ifdef BCH
    ss << "BN";
#endif 
  }
  ss << time(0) << ".out";
  return ss.str();
}

// write the performance results to a txt file.
void write_counters_to_file(){
  const string& filename = get_file_name();
#ifdef VERBOSE
  cout << "Results are saved to the following file: " << filename << endl;
#endif
  ofstream myfile (filename.c_str());
  if (myfile.is_open()){
    myfile << scrubbing_selection_switch << endl;
    myfile << trial_num << endl;
    myfile << fixed << setprecision(8);
    // write BERs
    for(long i = 0; i < (long) bit_error_rates.size(); i++){
      myfile << bit_error_rates[i] << " ";
    } 
    myfile << endl;
    myfile << average_iter_num  << endl;
    myfile << average_error_num << endl;
    myfile << average_correct_blk_num << endl;
    myfile << average_error_blk_num << endl;
    myfile << average_undet_blk_num << endl;
    myfile << average_det_blk_num << endl;
    myfile << average_final_ber << endl;
    myfile << average_mem_bit_num << endl;      
    myfile << average_cleanup_bit_num << endl;  
    myfile << num_seg_scrubbed << endl;

    myfile << average_ber_before_decode_meta << endl;
    myfile << average_iter_num_meta << endl;
    myfile << average_error_num_meta << endl;
    myfile << average_correct_blk_num_meta << endl;
    myfile << average_error_blk_num_meta << endl;
    myfile << average_compressed_metadata_size << endl;
    myfile.close();
  }
  else{
    cout << "The following data file can not be open:" << filename.c_str() << endl;
    exit(EXIT_FAILURE);
  }
}

// the main entry of the simulation.
int main(int argc, char *argv[]) {
  if(argc != 5){
    cout << "incorrect number of parameters" << endl;
    exit(EXIT_FAILURE);
  }

  // scheme selection 
  scrubbing_selection_switch = atoi(argv[1]);
  if ((scrubbing_selection_switch < 1) || (scrubbing_selection_switch > 2)) {
    cout << "The switch has to be either 1 or 2" << endl;
    exit(EXIT_FAILURE);
  }

  // total scrubbing time
  total_time = atof(argv[2]);
  // scrub gap 
  double gap = atof(argv[3]);
  bit_error_rates.resize((int) 2 * ((int) total_time)/gap);
  if (scrubbing_selection_switch == 1)
    scrub_gap = gap;
  else
    scrub_gap_naive = gap;  
  // number of trials
  trial_num = atoi(argv[4]);
  // initialize performance counters
  init_performance_counters();
#ifdef LDPC
  mat_GF2 H = mat_GF2();
  mat_GF2 G = mat_GF2();
  read_matrix_H(H, mat_dir_H1, M, N);
  LDPCEncoder::compute_G(G, H);
  if (DEBUG)
    LDPCEncoder::verify_G(G, H);
  simulate_scrub_ldpc(H, G);
#endif
#ifdef BCH
  simulate_scrub_bch();
#endif

  compute_average_performance();
  print_statistics();
  write_counters_to_file();
  return 0;
}
